////////////////////////////////////////////////
// DESCRIPTION:
//    Simple C/C++ Heap
//
// Legal Notices:
//     Copyright (c) 2008, Telliam Consulting, LLC.
//     All rights reserved.
//
//     Redistribution and use in source and binary forms, with or without modification,
//     are permitted provided that the following conditions are met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright notice,
//       this list of conditions and the following disclaimer in the documentation
//       and/or other materials provided with the distribution.
//
//     * Neither the name of Telliam Consulting nor the names of its contributors
//       may be used to endorse or promote products derived from this software
//       without specific prior written permission.
//
// Disclaimer:
//     THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
//     ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
//     WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
//     IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
//     INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
//     BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
//     DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
//     LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
//     OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
//     OF THE POSSIBILITY OF SUCH DAMAGE.  
//

//  Terms:
//    block    - A block is a piece of the heap that has a data format
//               that starts with the BlockHeader format and has zero or
//               more data bytes.
//
//    fragment - A fragment is the bytes of a block that are left when the
//               number of bytes allocated from that block is less than the
//               block data size.
//
//  Implementation Notes:
//    * There are two 'linked lists'
//      1) The linear list of adjacent blocks within the 'heap' array.
//         The size field serves as an offset to the next block in
//         the linear list.
//      2) The free space List which, will most likely not be linear in
//         relation to the underlying array. The free space "list"
//         has a head and tail node. (placed at the very begining and end of the
//         linear list ) The prevFree and nextFree fields for Allocated blocks
//         (tag=USED) are meaningless (except for head and tail).
//    * The data field of a block is where the "malloc'd" data starts.
//    * Minimum fragment block possible MUST be at least sizeof(BlockHeader)
//      bytes+1.
//    * Memory search is First Fit with adjustment. Adjacent Free blocks
//      are merged into a single large block.
//
#include <Windows.h>
#include <DBGAPI.h>

#if !defined(ERRFALSE)
#define ERRFALSE(exp)  extern char __ERRFALSECHECK[(exp)!=0]
#endif

typedef enum BlockTag
{
    USED,
    FREE
};

struct BlockHeader
{
    BlockHeader* prev;     // Previous Block (FREE or USED)
    unsigned size;         // Size of user data for this block
    BlockHeader* nextFree; // Next free block in list (meaningless if tag==USED)
    BlockHeader* prevFree; // Prev free block in list (meaningless if tag==USED)
    BlockTag tag;          // FREE or USED
};

// make certain the size is DWORD aligned
ERRFALSE(sizeof(BlockHeader) % sizeof(DWORD) == 0 );

void* operator new(size_t Size)
{
    return malloc(Size);
}

void operator delete(void* pMem)
{
    if(NULL != pMem)
        free(pMem);
}

/* next block from given block */
#define NEXT(x) ((BlockHeader *)(((unsigned char  *)(x)) + ((x)->size + sizeof(BlockHeader))))

#define minFragSize sizeof(BlockHeader) + sizeof(DWORD)     /* minimum frag block possible */

static int walkFreeListFwd( BlockHeader** where );
static int walkFreeListBwd( BlockHeader** where );
static int walkLinearListFwd( BlockHeader** where );
static int walkLinearListBwd( BlockHeader** where );
static int InHeap( void* p );
static int HeapCheck( void );
static void FreeList_addToHead( BlockHeader* blk );
static void FreeList_remove( BlockHeader* blk );
static void MergeBlockLeft( BlockHeader* blk );
static void ConstructBlock( BlockHeader* blk, BlockHeader* prevptr, BlockHeader* nextptr, size_t sz, int tag );
static void ShrinkBlock( BlockHeader* blk, size_t sz );

enum
{
    HEAPOK,            /* Heap is OK */
    BADHEADPOS,        /* Head indicated in wrong position */
    PREVNOTINHEAP,     /* prev pointer points to area not in heap */
    BADLINEARTAILPOS,  /* Tail found somewhere other than end of heap */
    BADFREEHEADPOS,    /* free list head indicated in wrong position */
    PREVFREENOTINHEAP, /* prevFree points to area not in heap */
    BADFREETAILPOS,    /* Free tail indicated in wrong position */ 
    NEXTFREENOTINHEAP, /* nextFree points to area not in heap */
    OUTOFMEMORY        /* out of memory Error! */ 
};

static TCHAR* HeapMsgs[] =
{
    _T("Heap OK\r\n"),
    _T("Head (prev==NULL) found at location!=s_pHeapStart\r\n"),
    _T("node prev points to area not in heap!\r\n"),
    _T("node found where NEXT()is outside heap and it is not the tail\r\n"),
    _T("In Back walk free node found with prevFree==NULL but location!=s_pFirstFree\r\n"),
    _T("node found where prevFree points to area outside of heap\r\n"),
    _T("In Forward walk free node found with nextFree==NULL but location!=real tail\r\n"),
    _T("node found where nextFree points to area outside of heap\r\n"),
    _T("Out of Memory in malloc()\r\n"),
};

/* STATIC DATA *****************************/
static unsigned char* s_pHeapStart;
static size_t s_HeapSize;
static BlockHeader* s_pFirstFree;

/***************************************
// DESCRIPTION:
//   Test to see if pointer points inside the heap
// INPUT:
//   void *p
// RETURN:
//   int  - 0 if not in heap; non zero if it is
// Comments:
//   Uses HUGE to force pointer normalization for comparison
//   in WIN32 all pointers are 32bit flat so they are by definition
//   normalized.
*/
inline int InHeap( void* p )
{
    return ( s_pHeapStart <= (unsigned char*)p && (unsigned char*)(p) <= s_pHeapStart + s_HeapSize - sizeof(BlockHeader) );
}

static void FullHeapInit(void* pHeapStart, size_t Size)
{
    BlockHeader* firstBlock;
    BlockHeader* tail;

    s_pFirstFree = ( BlockHeader* )pHeapStart;
    s_HeapSize = Size;
    s_pHeapStart = (BYTE*)pHeapStart;

    /* Create list head */
    s_pFirstFree->size = 0;
    s_pFirstFree->prev = NULL;
    s_pFirstFree->prevFree = NULL;
    s_pFirstFree->tag = USED;
    s_pFirstFree->nextFree = ( BlockHeader * )( s_pHeapStart + sizeof(BlockHeader) );

    /* initialize first Block */
    firstBlock = s_pFirstFree->nextFree;
    firstBlock->prev = s_pFirstFree;
    firstBlock->prevFree = s_pFirstFree;
    firstBlock->tag = FREE;
    firstBlock->size = (s_HeapSize - (sizeof(BlockHeader)*3));
    firstBlock->nextFree = ( BlockHeader* )( s_pHeapStart + ( s_HeapSize - sizeof(BlockHeader) ) );

    /* list tail */
    tail = firstBlock->nextFree;
    tail->prev = firstBlock;
    tail->prevFree = firstBlock;
    tail->nextFree = NULL;
    tail->tag = USED;
    tail->size = 0;
}

static void ReuseHeapInit(void* pHeapStart, size_t Size )
{
    s_pFirstFree = ( BlockHeader* )pHeapStart;
    s_HeapSize = Size;
    s_pHeapStart = (BYTE*)pHeapStart;

    ASSERT(HEAPOK == HeapCheck());
}

extern "C" {
    
    /***************************************
    // DESCRIPTION:
    //   Initialize the heap for first time use
    // Comments:
    //   Sets up the head and tail Blocks and establishes
    //   one large free block for allocation.
    */
    void InitHeap(void* pHeapStart, size_t Size, bool ForceFullInit )
    {
        if(ForceFullInit)
            FullHeapInit(pHeapStart, Size);
        else
            ReuseHeapInit(pHeapStart, Size);
    }

    /***************************************
    // DESCRIPTION:
    //   replacement for standard malloc()
    // INPUT:
    //   size_t size - size of memory block to allocate
    // RETURN:
    //   returns a pointer to the data block allocated
    //   or NULL on error.
    // Comments:
    //   Uses a first fit method of finding a free block
    //   Does not expand the heap when it is full (since it
    //   is a static block of data)
    */
    void* malloc( size_t size )
    {
        BlockHeader* pCurr;
    
        unsigned slack;
    
        ASSERT(HEAPOK == HeapCheck());

        pCurr = s_pFirstFree->nextFree;       /* Skip Free List head */
    
        while( pCurr->nextFree != NULL )      /* Find First Free block */
        {
            if( pCurr->size >= size )         /* Found a free block big enough ( not necessarily best fit) */
            {
                slack = pCurr->size - size;   /* how much bigger is it? */
            
                if( slack < minFragSize )     /* if 'slack' space is small enough */
                {
                    FreeList_remove( pCurr ); /* Just give 'em the whole thing and mark it as USED */
                    ASSERT(HEAPOK == HeapCheck());
                    return pCurr + 1;
                }
                ShrinkBlock( pCurr, size );    /* resize the block creating a new free one in the process */
                FreeList_remove( pCurr );      /* now remove it from the free list and mark as USED */
                ASSERT(HEAPOK==HeapCheck());
                return pCurr + 1;
            }
            pCurr = pCurr->nextFree;            /* Keep scaning through free list */
        }

        DEBUGMSG(1,( HeapMsgs[OUTOFMEMORY] ));
    
        ASSERT(HEAPOK == HeapCheck());
        return NULL;
    }

    /***************************************
    // DESCRIPTION:
    //  Replacement for standard realloc()
    // INPUT:
    //  void xdat *block - the user block to re allocate
    //  size_t size      - new size of block
    // RETURN:
    //  pointer to block with new size; not necessarily same as input.
    // Comments:
    //  Uses a close enough approach. If the difference is less than
    //  the Minimum fragmentation size then the block is left alone.
    */
    void* realloc( void* block, size_t size )
    {
        BlockHeader * curr;
    
        size_t fragSize;
        void*     newPtr;
    
        if( block == NULL )
            return malloc( size );

        ASSERT(HEAPOK==HeapCheck());
        curr = ( BlockHeader* )( ( ( BYTE* )block )-sizeof(BlockHeader) );
    
    
        /* back off pointer to get Block header */
        if( curr->size == size ) /* if same size just give it back to caller */
        {
            ASSERT(HEAPOK==HeapCheck());
            return block;
        }
    
        if( curr->size > size ) /* make it smaller */
        {
            fragSize = curr->size - size;    /* how much smaller */
        
            if( fragSize >= minFragSize )    /* is it worth the effort? */
                ShrinkBlock( curr, size );    /* shrink it to new size creating new free one in the process */
            ASSERT(HEAPOK==HeapCheck());
            return block;    /* hand back the given pointer */
        }
    
        /* ... else make it bigger */
        newPtr = malloc( size );    /* malloc a new block then free the old one */
    
        if( newPtr != NULL )
        {
            memcpy( newPtr, block, curr->size );    /* copy data area */
            free( block );
        }
        ASSERT(HEAPOK==HeapCheck());
        return newPtr;
    }

    /***************************************
    // DESCRIPTION:
    //   Replacement for Standard free()
    // INPUT:
    //   void  *p - pointer to user block to free
    // Comments:
    //   Places the block back into the free list merging it with
    //   any neighboring free blocks.
    //   Right and left references in this context are to
    //   the linear list within the heap! NOT the free List.
    */
    void free( void* p )
    {
        BlockHeader* curr, *leftN, *rightN;
        
        ASSERT(HEAPOK==HeapCheck());
        ASSERT(InHeap(p));

        curr = ((( BlockHeader* )p ) -1 );
        leftN = curr->prev;
        rightN = NEXT( curr );
    
        if( leftN->tag == FREE && rightN->tag == FREE ) /* both neighbors free */
        {
            MergeBlockLeft( rightN );    /* merge right one into this one */
            MergeBlockLeft( curr );    /* merge this one into Left one */
        }
        else if( rightN->tag == FREE ) /* right neighbor free */
        {
            MergeBlockLeft( rightN );    /* merge the right one into this one */
            FreeList_addToHead( curr );    /* now add it to the free List */
        }
        else if( leftN->tag == FREE )    /* left neighbor free */
            MergeBlockLeft( curr );    /* merge this one into left neighbor */
        else
            FreeList_addToHead( curr );    /* no neighbors free (stick at head of free space list)*/

        ASSERT(HEAPOK==HeapCheck());
    }

    /***************************************
    // DESCRIPTION:
    //   Walk Free list forward Adding up block sizes
    // RETURN:
    //   size_t - size of all free blocks (not including Overhead)
    */
    FreeMemSize()
    {
        BlockHeader * curr;
    
        size_t freemem = 0;
    
        curr = s_pFirstFree;
    
        while( curr->nextFree != NULL )
        {
            freemem += curr->size;
            curr = curr->nextFree;
        }
        return freemem;
    }
}




/***************************************
// DESCRIPTION:
//   Put a block at the head of the FREE list.
// INPUT:
//   BlockHeader *blk - pointer to the block
*/
static void FreeList_addToHead( BlockHeader* blk )
{
    blk->tag = FREE;    /* mark the block as FREE */
    blk->nextFree = s_pFirstFree->nextFree;    /* link this one forward to next */
    blk->prevFree = s_pFirstFree;    /* link this one back to list head */
    blk->nextFree->prevFree = blk;    /* link next one back to this one */
    s_pFirstFree->nextFree = blk;    /* link list head forward to this one*/
}

/***************************************
// DESCRIPTION:
//   Remove Block from the free list.
// INPUT:
//   BlockHeader *blk
// Comments:
//   Only reason to remove from free list is as
//   part of malloc so the block is marked as used.
*/
static void FreeList_remove( BlockHeader* blk )
{
    blk->tag = USED;    /* only reason to remove is to make it used */
    blk->prevFree->nextFree = blk->nextFree;    /* link previous one around to next one */
    blk->nextFree->prevFree = blk->prevFree;    /* link next one back around to previous one */
}

/***************************************
// DESCRIPTION:
//   Merge a block into the block to it's left
// INPUT:
//   BlockHeader *blk - the block to merge with it's left neighbor
// Comments:
//   Left in this context refers to position in the Linear List.
*/
static void MergeBlockLeft( BlockHeader* blk )
{
    if( blk->tag == FREE )
        FreeList_remove( blk );    /* pull it out of free list */

    blk->prev->size += blk->size + sizeof(BlockHeader);    /* increase size of left one to consume this one */
    NEXT( blk )->prev = blk->prev;    /* link next one around to previous */
}

/***************************************
// DESCRIPTION:
//   Initialize data members of a block
// INPUT:
//   BlockHeader *self     - Block to be initialized
//   BlockHeader *prevptr  - init value for new block
//   BlockHeader *nextptr  - init value for new block
//   size_t sz       - init value for new block
//   int tag         - init state for new block (USED,FREE)
*/
static void ConstructBlock( BlockHeader* self, BlockHeader* prevptr, BlockHeader* nextptr, size_t sz, BlockTag tag )
{
    self->prev = prevptr;    /* link this one back to the previous one */
    self->size = sz;
    nextptr->prev = self;    /* link the next one back to this one */
    self->tag = tag;
}

/***************************************
// DESCRIPTION:
//   Reduce the size of a block
// INPUT:
//    BlockHeader *self - the block to shrink
//    size_t sz   - new size of block
// Comments:
//    sz is assumed to be less than the current size of
//    the block. If it is not un predictable results
//    will occur. Most likely leading to data corruption.
*/
static void ShrinkBlock( BlockHeader* self, size_t sz )
{
    BlockHeader* temp;
    temp = ( BlockHeader * )( ( ( unsigned char* )self )+( sz + sizeof(BlockHeader) ) );
    ConstructBlock( temp, self, NEXT( self ), self->size- sz- sizeof(BlockHeader), FREE );
    
    /* construct new block with slack as size */
    self->size = sz;
    FreeList_addToHead( temp );    /* add new block to free List */
}

/***************************************
// DESCRIPTION:
//   Walk Free list forward Testing for errors
// INPUT:
//   BlockHeader **where - [OUT] place to put pointer to errant block
// RETURN:
//   int - heap error code (see list above)
// Comments:
//   Walks through the list to make sure head is consistent
//   with tail and that no pointer point outside of the actual
//   heap area.
*/
static int walkFreeListFwd( BlockHeader** where )
{
    BlockHeader * curr;
    curr = s_pFirstFree;
    
    while( curr->nextFree != NULL )
    {
        if( !InHeap( curr->nextFree ) )
        {
            *where = curr;
            return NEXTFREENOTINHEAP;    /* next free points outside of heap */
        }
        else
            curr = curr->nextFree;
    }
    
    if( curr != ( BlockHeader * )( s_pHeapStart + s_HeapSize - sizeof(BlockHeader) ) ) /* is it really the tail? */
    {
        *where = curr;
        return BADFREETAILPOS;    /* free tail indicated where it isn't supposed to be */
    }
    return HEAPOK;
}




/***************************************
// DESCRIPTION:
//   Walk Free list backward Testing for errors
// INPUT:
//   BlockHeader **where - [OUT] place to put pointer to errant block
// RETURN:
//   int - heap error code (see list above)
// Comments:
//   Walks through the list to make sure head is consistent
//   with tail and that no pointer point outside of the actual
//   heap area.
*/
static int walkFreeListBwd( BlockHeader** where )
{
    BlockHeader * curr;
    curr = ( BlockHeader * )( s_pHeapStart + ( s_HeapSize - sizeof(BlockHeader) ) );
    
    while( curr->prevFree != NULL )
    {
        if( !InHeap( curr->prevFree ) )
        {
            *where = curr;
            return PREVFREENOTINHEAP;    /* prev free points outside of heap */
        }
        else
            curr = curr->prevFree;
    }
    
    if( curr != s_pFirstFree ) /* is it really the head */
    {
        *where = curr;
        return BADFREEHEADPOS;    /* free head indicated where it isn't supposed to be */
    }
    return HEAPOK;
}




/***************************************
// DESCRIPTION:
//   Walk Linear list forward Testing for errors
// INPUT:
//   BlockHeader **where - [OUT] place to put pointer to errant block
// RETURN:
//   int - heap error code (see list above)
// Comments:
//   Walks through the list to make sure head is consistent
//   with tail and that no pointer point outside of the actual
//   heap area.
*/
static int walkLinearListFwd( BlockHeader** where )
{
    BlockHeader * curr;
    curr = ( BlockHeader * )s_pHeapStart;
    
    while( InHeap( NEXT( curr ) ) )
        curr = NEXT( curr );
    
    if( curr != ( BlockHeader * )( s_pHeapStart + s_HeapSize - sizeof(BlockHeader) ) ) /* is it really the tail? */
    {
        *where = curr;
        return BADLINEARTAILPOS;    /* Linear tail indicated where it isn't supposed to be */
    }
    return HEAPOK;
}




/***************************************
// DESCRIPTION:
//   Walk Linear list backward Testing for errors
// INPUT:
//   BlockHeader **where - [OUT] place to put pointer to errant block
// RETURN:
//   int - heap error code (see list above)
// Comments:
//   Walks through the list to make sure head is consistent
//   with tail and that no pointer point outside of the actual
//   heap area.
*/
static int walkLinearListBwd( BlockHeader** where )
{
    BlockHeader * curr;
    curr = ( BlockHeader * )( s_pHeapStart + ( s_HeapSize - sizeof(BlockHeader) ) );
    
    while( curr->prev != NULL )
    {
        if( !InHeap( curr->prev ) )
        {
            *where = curr;
            return PREVNOTINHEAP;    /* prev points outside of heap */
        }
        else
            curr = curr->prev;
    }
    
    if( curr != s_pFirstFree ) /* is it really the head */
    {
        *where = curr;
        return BADHEADPOS;    /* Linear head indicated where it isn't supposed to be */
    }
    return HEAPOK;
}




/***************************************
// DESCRIPTION:
//   Check the heap for consistency
// Comments:
//   Prints out error messages when
//   inconsistencies are detected.
*/
int HeapCheck( void )
{
    BlockHeader* loc;
    int err;

    err = walkFreeListFwd( &loc );
    DEBUGMSG(err != HEAPOK,( HeapMsgs[err] ));
    
    err = walkFreeListBwd( &loc );
    DEBUGMSG(err != HEAPOK,( HeapMsgs[err] ));

    err = walkLinearListFwd( &loc );
    DEBUGMSG(err != HEAPOK,( HeapMsgs[err] ));

    err = walkLinearListBwd( &loc );
    DEBUGMSG(err != HEAPOK,( HeapMsgs[err] ));

    return err;
}


